abstract type Block end

function steady_state(b::Block, calibration; dissolve=[], options=Dict(), kwargs...)
    inputs = copy(b.inputs)

    if b isa ParentBlock
        for k in dissolve
            inputs = inputs ∪ keys(get_attribute(b, k, :unknowns))
        end
    end

    calibration = SteadyStateDict(calibration)[inputs]
    own_options = get_options(b, options, kwargs, "steady_state")

    if b isa ParentBlock
        return b.M * _steady_state(b, inv(b.M) * calibration; dissolve=dissolve, options=options, own_options...)
    else
        return b.M * _steady_state(b, inv(b.M) * calibration; own_options...)
    end
end

function solve_for_unknowns(f!, unknowns, solver, solver_kwargs; kwargs...)
    NLsolve.nlsolve(f!, collect(values(unknowns)), show_trace=true, iterations=2000)
end

function solve_steady_state(b::Block, calibration, unknowns, targets; dissolve = [], options = Dict(), kwargs...)
    options = get_options(b, options, kwargs, "solve_steady_state")

    ss = SteadyStateDict(calibration)

    solver = !isempty(options["solver"]) ? options["solver"] : provide_solver_default(unknowns)

    function residual!(F::AbstractArray, unknown_values::AbstractArray; unknowns_keys=keys(unknowns), targets=targets)
        merge!(ss, zip(unknowns_keys, unknown_values))
        merge!(ss, steady_state(b, ss; dissolve=dissolve, options=options, kwargs...))
        for i ∈ collect(1:length(F))
            F[i] = compute_target_values(targets, ss)[i]
        end
        return nothing
    end
    
    sol = solve_for_unknowns(residual!, unknowns, solver, options["solver_kwargs"];
            tol=options["ttol"], verbose=options["verbose"],
            constrained_method=options["constrained_method"],
            constrained_kwargs=options["constrained_kwargs"])

    residual!(zeros(length(unknowns)), sol.zero)

    return ss
end

function impulse_linear(b::Block, ss, inputs, outputs, Js, options, kwargs...)
end

function solve_impulse_linear(b::Block, ss, unknowns, targets, inputs, outputs, Js, options, H_U_factored, kwargs...)
end

function partial_jacobians(b::Block, ss, inputs = nothing, outputs = nothing; T = nothing, Js = Dict(), options = Dict(), kwargs...)
    if inputs isa Nothing
        inputs = b.inputs
    end
    if outputs isa Nothing
        outputs = b.outputs
    end

    if (b.name ∈ keys(Js)) && (Js[b.name] isa JacobianDict) && (inputs ⊆ Js[b.name].inputs) && (outputs ⊆ Js[b.name].outputs)
        return Dict(b.name => Js[b.name][outputs, inputs])
    end

    if !(b isa ParentBlock)
        own_options = get_options(b, options, kwargs, "jacobian")
        jac = jacobian(b, ss, inputs, outputs; T=T, own_options...)
        return !isempty(jac) ? Dict(b.name => jac) : Dict()
    end

    own_options = get_options(b, options, kwargs, "partial_jacobians")

    partial = _partial_jacobians(b, inv(b.M) * ss, inv(b.M) * inputs, inv(b.M) * outputs; T=T, Js=Js, options=options, own_options...)

    if b.name ∈ keys(partial)
        partial[b.name] = b.M * partial[b.name]
    end
    return partial
end

function jacobian(b::Block, ss, inputs, outputs = nothing; T = nothing, Js = Dict(), options = Dict(), kwargs...)
    own_options = get_options(b, options, kwargs, "jacobian")
    inputs = make_ordered_set(inputs)
    outputs, _ = process_outputs(b, ss, Dict(), make_ordered_set(outputs))

    if b.name ∈ keys(Js) && Js[b.name] isa JacobianDict && inputs ⊆ Js[b.name].inputs && outputs ⊆ Js[b.name].outputs
        return Js[b.name][outputs, inputs]
    end

    if !(b isa ParentBlock)
        return b.M * _jacobian(b, inv(b.M) * ss, inv(b.M) * inputs, inv(b.M) * outputs; T=T, own_options...)
    end

    Js = copy(Js)
    if b.name ∈ keys(Js)
        Js[b.name] = inv(b.M) * Js[b.name]
    end
    return b.M * _jacobian(b, inv(b.M) * ss, inv(b.M) * inputs, inv(b.M) * outputs; T=T, Js=Js, options=options, own_options...)
end

Base.union(::Nothing, x::Any) = x
Base.union(x::Any, ::Nothing) = x

function solve_jacobian(b::Block, ss, unknowns, targets, inputs, outputs = nothing; T=300, Js = Dict(), options = Dict(), H_U_factored = nothing, kwargs...)
    inputs, unknowns = make_ordered_set(inputs), make_ordered_set(unknowns)
    actual_outputs, unknowns_as_outputs = process_outputs(b, ss, unknowns, make_ordered_set(outputs))

    Js = partial_jacobians(b, ss, inputs ∪ unknowns, setdiff(actual_outputs ∪ targets, unknowns); T=T, Js=Js, options=options, kwargs...)

    H_Z = jacobian(b, ss, inputs, targets; T=T, Js=Js, options=options, kwargs...)

    if H_U_factored isa Nothing
        H_U = pack(jacobian(b, ss, unknowns, targets; T=T, Js=Js, options=options, kwargs...); T=T)
        U_Z = unpack(- H_U \ pack(H_Z; T=T), unknowns, inputs, T)
    else
        U_Z = H_U_factored * H_Z
    end

    self_with_unknowns = combine([U_Z, b])

    return jacobian(self_with_unknowns, ss, inputs, unknowns_as_outputs ∪ actual_outputs; T=T, Js=Js, options=options, kwargs...)
end

function impulse_nonlinear(b::Block, ss, inputs; outputs = nothing, internals = Dict(), Js = Dict(), options = Dict(), ss_initial = nothing, kwargs... )
    own_options = get_options(b, options, kwargs, "impulse_nonlinear")
    inputs = ImpulseDict(inputs)
    actual_outputs, inputs_as_outputs = process_outputs(b, ss, make_ordered_set(inputs), make_ordered_set(outputs))

    if b isa ParentBlock
        out = b.M * _impulse_nonlinear(b, inv(b.M) * ss, inv(b.M) * inputs; outputs = inv(b.M) * actual_outputs, internals=internals, Js=Js, options=options, ss_initial = inv(b.M)* ss_initial, own_options...)
    elseif hasproperty(b, :internals)
        out = b.M * _impulse_nonlinear(b, inv(b.M) * ss, inv(b.M) * inputs; outputs = inv(b.M) * actual_outputs, internals = internals_to_report(b, internals), ss_initial = inv(b.M)* ss_initial, own_options...)
    else
        out = b.M * _impulse_nonlinear(b, inv(b.M) * ss, inv(b.M) * inputs; outputs = inv(b.M) * actual_outputs, ss_initial = inv(b.M) * ss_initial, own_options...)
    end
    return inputs[inputs_as_outputs] ∪ out
end

function solve_impulse_nonlinear(b::Block, ss, unknowns, targets, inputs; outputs = nothing, internals = Dict(), Js = Dict(), options = Dict(), H_U_factored = nothing, ss_initial = nothing, kwargs...)
    inputs  = ImpulseDict(inputs)
    unknowns, targets = OrderedSet(unknowns), OrderedSet(targets)

    input_names = make_ordered_set(keys(inputs))
    actual_outputs, inputs_as_outputs = process_outputs(b, ss, union(input_names, unknowns), make_ordered_set(outputs))

    T = inputs.T

    Js = partial_jacobians(b, ss, union(input_names, unknowns), setdiff(union(actual_outputs, targets), unknowns); T=T, Js=Js, options=options, kwargs...)

    if isnothing(H_U_factored)
        H_U  = jacobian(b, ss, unknowns, targets; T=T, Js=Js, options=options, kwargs...)
        H_U_factored = FactoredJacobianDict(H_U; T=T)
    end

    options = get_options(b, options, kwargs, "solve_impulse_nonlinear")

    # Use Newton's Method to iterate for nonlinear impulse responses
    U = ImpulseDict(Dict(k => zeros(T) for k ∈ unknowns))
    if options["verbose"]
        println("Solving '$(b.name)' for '$(unknowns)' to hit '$(targets)'")
    end
    results = ImpulseDict(Dict(k => zeros(T) for k ∈ unknowns)) # results casted as array of zeros
    does_break = false
    for it ∈ 1:options["maxit"] # iterate over loop until convergence
        results = impulse_nonlinear(b, ss, inputs ∪ U; outputs = actual_outputs ∪ targets, internals=internals, Js=Js, options=options, ss_initial=ss_initial, kwargs...) # computes results for endogenous inputs
        errors = Dict(k => maximum(abs.(results[k])) for k ∈ targets) # computes max errors for goods market and euler equations for each iteration; k in targets are goods_mkt and euler
        if options["verbose"]
            println("On iteration $it") # gives current iteration results
            for k ∈ errors
                println("    max error for $k is $errors[k]") # gives max errors for each iteration of goods market and euler
            end
        end

        if all(v .< options["tol"] for v ∈ values(errors))
            does_break = true
            break
        else
            U += apply(H_U_factored, results) # updates for each iteration
        end
    end
    if !does_break
        error("No convergence after $(options["maxit"]) backward iterations!")
    end
    return (inputs ∪ U)[inputs_as_outputs] ∪ results
end

function solved(b::Block, unknowns, targets; name=nothing, solver=nothing, solver_kwargs=nothing)
    isnothing(name) && (name = b.name * "_solved")
    return SolvedBlock(b, name, unknowns, targets; solver=solver, solver_kwargs=solver_kwargs)
end

function remap(b::Block, map)
    other = deepcopy(b)
    other.M = b.M * Bijection(map)
    other.inputs = other.M * b.inputs
    other.outputs = other.M * b.outputs
    if hasproperty(b, :input_list)
        other.input_list = other.M * b.input_list
    end
    if hasproperty(b, :output_list)
        other.output_list = other.M * b.output_list
    end
    if hasproperty(b, :output_list)
        other.non_back_iter_outputs = other.M * b.non_back_iter_outputs
    end
    return other
end

function rename(b::Block; name=nothing, suffix=nothing)
    if b isa ParentBlock
        other = deepcopy(b)
        other.name = b.name * suffix
        if hasproperty(b, :blocks)
            other.blocks = [block.rename(name, suffix) for block ∈ b.blocks]
            other.name, other.kids, other.descendants = construct_parent(other.blocks; name=other.name)
        elseif hasproperty(b, :block)
            other.block = rename_top(b.block, b.block.name * suffix)
            other.name, other.kids, other.descendants = construct_parent([other.block]; name=other.name)
        end
        return other
    else
        if isnothing(suffix)
            return rename_top(b, name)
        else
            return rename_top(b, b.name * suffix)
        end
    end
end

function rename_top(b::Block, name)
    other = deepcopy(b)
    other.name = name
    return other
end

function default_inputs_outputs(b::Block, ss, inputs, outputs)
    if inputs isa Nothing
        inputs = b.inputs
    end
    if outputs isa Nothing
        outputs = setdiff(b.outputs, ss._vector_valued())
    end
    return OrderedSet(inputs), OrderedSet(outputs)
end

function process_outputs(b::Block, ss, inputs, outputs)
    if outputs isa Nothing
        actual_outputs = setdiff(b.outputs, _vector_valued(ss))
        inputs_as_outputs = inputs
    else
        actual_outputs = outputs ∩ b.outputs
        inputs_as_outputs = outputs ∩ inputs
    end

    return actual_outputs, inputs_as_outputs
end

make_ordered_set(x::Union{Nothing, OrderedSet}) = x
make_ordered_set(x::AbstractDict) = OrderedSet(collect(keys(x)))
make_ordered_set(x) = OrderedSet(x)

macro add_options(method)
    return :()
end

function get_options(b::Block, options, kwargs, method)
    # is this okay?? no
    own_options = getproperty(b, Symbol(method * "_options"))

    str_kwargs = Dict(string(k) => v for (k, v) ∈ kwargs)

    if b.name ∈ keys(options)
        merged = Dict(own_options..., options[b.name]..., str_kwargs...)
    else
        merged = Dict(own_options..., str_kwargs...)
    end

    return Dict(k => merged[k] for k ∈ keys(own_options))
end

function input_defaults_smart(b::Block, methodname)
    method = hasproperty(b, methodname) ? getproperty(b, methodname) : nothing
    if method isa Nothing
        return Dict()
    else
        return input_defaults(method)
    end
end

function internals_to_report(b::Block, internals)
    if b.name ∈ internals
        if internals isa AbstractDict
            return internals[b.name]
        else
            return b.internals
        end
    else
        return []
    end
end
